package com.feidee.fd.sml.algorithm.ml

import com.feidee.fd.sml.algorithm.component.ml.classification.{DecisionTreeComponent, DecisionTreeParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.scalatest.FunSuite


/**
  * @Author JunxinWang
  * @Date 2019/3/25 15:53
  * @Description
  * @Reviewer
  */
class DecisionTreeComponentSuite extends FunSuite {

  val dt = new DecisionTreeComponent
  val paramStr: String =
    """
      |{
      |	'input_pt': 'a',
      |	'output_pt': 'a',
      |	'featuresCol': 'features',
      |	'labelCol': 'label',
      | 'probabilityCol': 'probability',
      |	'maxDepth': 5,
      |	'modelPath': '',
      | 'metrics': ['accuracy', 'auc', 'pr', 'abc']
      |}
    """.stripMargin
  val param: DecisionTreeParam = dt.parseParam(new ToolClass().encrypt(paramStr))

  test("decision tree parameter construction test") {
    assert("a".equals(param.input_pt) & "a".equals(param.output_pt) & "features".equals(param.featuresCol) &
      "label".equals(param.labelCol) & "rawPrediction".equals(param.rawPredictionCol) & "prediction".equals(param.predictionCol) &
      "probability".equals(param.probabilityCol) & "".equals(param.modelPath) & param.thresholds.length == 0
    )
  }

  test("decision tree training function test") {
    // 测试模型训练方法
    param.verify()
    val testingData = TestingDataGenerator.makeOrderedTrainingData(5, 10)
    val newData = dt.process(param, testingData)
    testingData.show()
    newData.show()
    //    assert(result._1.numClasses == 2)
    //    assert(result._2.count() == 20)
  }

}
